# -*- coding: utf-8 -*-
import numpy as np
import pandas as pd
import pickle
import tensorflow as tf
from tensorflow.keras import layers 
from tensorflow.keras.layers import Dense, LSTM
import matplotlib.pyplot as plt
from sklearn.metrics import mean_squared_error, make_scorer

font = {'family' : 'Times New Roman',
'weight' : 'normal',
'size'   : 20,
}
plt.rcParams['font.sans-serif']=['SimHei']
plt.rcParams['axes.unicode_minus'] = False


def fit_lstm_model(X, Y, fit_flag=True):
    lstm_model = tf.keras.models.Sequential([
        tf.keras.layers.LSTM(64, return_sequences=True),
        tf.keras.layers.Dense(units=7, activation='linear')
    ])
    
    optimizer = tf.keras.optimizers.Adam(
                                learning_rate=0.001, beta_1=0.9, beta_2=0.999, epsilon=1e-07, amsgrad=False,
                                            )
    checkpoint_filepath = r'../附件/lstm_model'
    model_checkpoint_callback = tf.keras.callbacks.ModelCheckpoint(
                                            filepath=checkpoint_filepath,
                                            save_weights_only=True,
                                            verbose=True,
                                            monitor='loss',
                                            mode='min',
                                            save_best_only=True)
    
    lstm_model.compile(loss='mse', optimizer=optimizer)
    #callback_list=[model_checkpoint_callback] 
    
    if fit_flag:
        #history = lstm_model.fit(X, Y, epochs=250, callbacks=callback_list)
        history = lstm_model.fit(X, Y, epochs=500)
    #lstm_model.load_weights(checkpoint_filepath)
    return history, lstm_model

def plot_result(history):
    '''
    会出迭代历史
    '''
    data = history.history['loss']
    fig = plt.figure(figsize=(6,4))

    plt.plot(data, linewidth=3)
    plt.xlabel('Epoch',fontsize=20)
    plt.ylabel('MSE',fontsize=20)
    plt.title('模型训练过程')
    fig.savefig('../图片/LSTM迭代历史.png')
    plt.show()

def plot_fitting(Y, Y_pred):
    for i in range(7):
        y = Y[:, i]
        y_pred = Y_pred[:, i]
        x = range(2,len(y)+2)
                
        fig = plt.figure()
        plt.plot(x, y, label='实际数据')
        plt.plot(x, y_pred, label='拟合数据')
        plt.xlabel('月份')
        plt.ylabel('发生次数')
        plt.title(f'类别 {i+1} 发生的次数（月）')
        fig.savefig(f'../图片/LSTM模型拟合数据与实际数据_时间类别{i+1}')
        plt.show()

if  __name__ == '__main__':
    data = pd.read_excel(r'../附件/各类事件（滑动窗口后）月次数数据.xlsx', index_col=0)
    # 前一天的数据用于预测
    X = data.iloc[:, :7].values.reshape((11, 1, 7))
    # 当天数据被预测。。
    Y = data.iloc[:, 7:].values.reshape((11, 1, 7))
    
    scaler = pickle.load(open(r'../附件/scaler_q3.pkl', 'rb'))
    
    history, lstm_model = fit_lstm_model(X, Y)
    # 模型预测
    Y_pred = lstm_model.predict(X)
    # MSE
    Y = scaler.inverse_transform(Y.reshape((11,7)))
    Y_pred = np.round(scaler.inverse_transform(Y_pred.reshape((11,7))))
    plot_result(history)
    # 均方误差（标准化前）
    score = mean_squared_error(Y, Y_pred)
    print('LSTMs 模型的拟合数据与数据数据间的均方差： ', score)
    plot_fitting(Y, Y_pred)
